"""
ollama_llm.py 模块
基于Ollama的LLM实现
"""
from typing import List, Optional, Dict
import ollama
from .base_llm import BaseLLM

class OllamaLLM(BaseLLM):
    """基于Ollama的LLM实现"""

    def __init__(
        self,
        model_name: str,
        temperature: float = 0.7,
        max_tokens: int = 1024,
        **kwargs
    ):
        """
        初始化Ollama LLM

        Args:
            model_name: 模型名称（如"qwen:7b"、"llama3:7b"等）
            temperature: 温度参数，控制生成文本的随机性
            max_tokens: 生成文本的最大长度
            **kwargs: 其他参数
        """
        super().__init__(model_name, temperature, max_tokens,** kwargs)

        # 检查模型是否可用
        try:
            ollama.show(model_name)
        except Exception as e:
            raise ValueError(f"Ollama模型 {model_name} 不可用，请确保已安装并拉取该模型。错误: {str(e)}")

    def generate(
        self,
        prompt: str,
        stop: Optional[List[str]] = None
    ) -> str:
        """
        生成文本

        Args:
            prompt: 提示词
            stop: 停止词列表，遇到这些词则停止生成

        Returns:
            生成的文本
        """
        response = ollama.generate(
            model=self.model_name,
            prompt=prompt,
            options={
                "temperature": self.temperature,
                "num_predict": self.max_tokens,
                **self.kwargs
            },
            stop=stop
        )
        return response["response"]

    def chat(
        self,
        messages: List[Dict[str, str]],
        stop: Optional[List[str]] = None
    ) -> str:
        """
        聊天接口，支持多轮对话

        Args:
            messages: 消息列表，每个消息是一个字典，包含"role"和"content"
            stop: 停止词列表，遇到这些词则停止生成

        Returns:
            生成的回复文本
        """
        response = ollama.chat(
            model=self.model_name,
            messages=messages,
            options={
                "temperature": self.temperature,
                "num_predict": self.max_tokens,** self.kwargs
            },
            stop=stop
        )
        return response["message"]["content"]
